package dns

import (
	"context"
	"fmt"
	"net"
	"os/exec"
	"strings"
	"time"
)

// RecordType DNS记录类型
type RecordType string

const (
	TypeA     RecordType = "A"
	TypeAAAA  RecordType = "AAAA"
	TypeCNAME RecordType = "CNAME"
	TypeMX    RecordType = "MX"
	TypeTXT   RecordType = "TXT"
	TypeNS    RecordType = "NS"
	TypeCAA   RecordType = "CAA"
	TypeSOA   RecordType = "SOA"
	TypePTR   RecordType = "PTR"
	TypeSRV   RecordType = "SRV"
)

// DNSRecord DNS记录
type DNSRecord struct {
	Name     string     `json:"name"`
	Type     RecordType `json:"type"`
	TTL      int        `json:"ttl"`
	Value    string     `json:"value"`
	Priority uint16     `json:"priority,omitempty"` // For MX records
	Weight   uint16     `json:"weight,omitempty"`   // For SRV records
	Port     uint16     `json:"port,omitempty"`     // For SRV records
}

// CAARecord CAA记录
type CAARecord struct {
	Domain string `json:"domain"`
	Flags  uint8  `json:"flags"`
	Tag    string `json:"tag"`
	Value  string `json:"value"`
}

// QueryResult 查询结果
type QueryResult struct {
	Domain    string        `json:"domain"`
	A         []DNSRecord   `json:"a,omitempty"`
	AAAA      []DNSRecord   `json:"aaaa,omitempty"`
	CNAME     []DNSRecord   `json:"cname,omitempty"`
	MX        []DNSRecord   `json:"mx,omitempty"`
	TXT       []DNSRecord   `json:"txt,omitempty"`
	NS        []DNSRecord   `json:"ns,omitempty"`
	CAA       []CAARecord   `json:"caa,omitempty"`
	QueryTime time.Duration `json:"query_time"`
}

// LocalResolver 本地DNS解析器
type LocalResolver struct {
	servers  []string
	resolver *net.Resolver
	timeout  time.Duration
}

// NewResolver 创建DNS解析器
func NewResolver(servers []string) *LocalResolver {
	r := &LocalResolver{
		servers: servers,
		timeout: 5 * time.Second,
	}

	if len(servers) > 0 {
		r.resolver = &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{
					Timeout: r.timeout,
				}
				// Use first available server
				for _, server := range servers {
					conn, err := d.DialContext(ctx, network, server+":53")
					if err == nil {
						return conn, nil
					}
				}
				// Fallback to default
				return d.DialContext(ctx, network, address)
			},
		}
	} else {
		r.resolver = net.DefaultResolver
	}

	return r
}

// SetTimeout 设置超时时间
func (r *LocalResolver) SetTimeout(timeout time.Duration) {
	r.timeout = timeout
}

// Query 查询DNS记录
func (r *LocalResolver) Query(ctx context.Context, domain string, recordType RecordType) ([]DNSRecord, error) {
	if domain == "" {
		return nil, fmt.Errorf("domain cannot be empty")
	}

	// Remove trailing dot if present
	domain = strings.TrimSuffix(domain, ".")

	var records []DNSRecord

	switch recordType {
	case TypeA:
		ips, err := r.resolver.LookupIPAddr(ctx, domain)
		if err != nil {
			return nil, err
		}
		for _, ip := range ips {
			if ip.IP.To4() != nil {
				records = append(records, DNSRecord{
					Name:  domain,
					Type:  TypeA,
					TTL:   300,
					Value: ip.IP.String(),
				})
			}
		}

	case TypeAAAA:
		ips, err := r.resolver.LookupIPAddr(ctx, domain)
		if err != nil {
			return nil, err
		}
		for _, ip := range ips {
			if ip.IP.To4() == nil && ip.IP.To16() != nil {
				records = append(records, DNSRecord{
					Name:  domain,
					Type:  TypeAAAA,
					TTL:   300,
					Value: ip.IP.String(),
				})
			}
		}

	case TypeCNAME:
		cname, err := r.resolver.LookupCNAME(ctx, domain)
		if err != nil {
			return nil, err
		}
		if cname != domain && cname != domain+"." {
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeCNAME,
				TTL:   300,
				Value: strings.TrimSuffix(cname, "."),
			})
		}

	case TypeMX:
		mxs, err := r.resolver.LookupMX(ctx, domain)
		if err != nil {
			return nil, err
		}
		for _, mx := range mxs {
			records = append(records, DNSRecord{
				Name:     domain,
				Type:     TypeMX,
				TTL:      300,
				Value:    strings.TrimSuffix(mx.Host, "."),
				Priority: mx.Pref,
			})
		}

	case TypeTXT:
		txts, err := r.resolver.LookupTXT(ctx, domain)
		if err != nil {
			return nil, err
		}
		for _, txt := range txts {
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeTXT,
				TTL:   300,
				Value: txt,
			})
		}

	case TypeNS:
		nss, err := r.resolver.LookupNS(ctx, domain)
		if err != nil {
			return nil, err
		}
		for _, ns := range nss {
			records = append(records, DNSRecord{
				Name:  domain,
				Type:  TypeNS,
				TTL:   300,
				Value: strings.TrimSuffix(ns.Host, "."),
			})
		}

	default:
		return nil, fmt.Errorf("unsupported record type: %s", recordType)
	}

	return records, nil
}

// QueryCAA 查询CAA记录
func (r *LocalResolver) QueryCAA(ctx context.Context, domain string) ([]CAARecord, error) {
	if domain == "" {
		return nil, fmt.Errorf("domain cannot be empty")
	}

	// Remove trailing dot
	domain = strings.TrimSuffix(domain, ".")

	var records []CAARecord

	// Try using dig command for CAA records
	args := []string{"+short", "+tcp", "CAA", domain}
	if len(r.servers) > 0 && r.servers[0] != "" {
		args = append([]string{"@" + r.servers[0]}, args...)
	}

	cmd := exec.CommandContext(ctx, "dig", args...)
	output, err := cmd.Output()
	if err != nil {
		// If dig fails, return empty records (CAA might not be supported)
		return records, nil
	}

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.Contains(line, "connection timed out") {
			continue
		}

		// Parse CAA record format: "flags tag value"
		// Example: 0 issue "letsencrypt.org"
		parts := strings.Fields(line)
		if len(parts) >= 3 {
			var flags uint8
			fmt.Sscanf(parts[0], "%d", &flags)

			tag := parts[1]
			value := strings.Join(parts[2:], " ")
			value = strings.Trim(value, "\"")

			records = append(records, CAARecord{
				Domain: domain,
				Flags:  flags,
				Tag:    tag,
				Value:  value,
			})
		}
	}

	return records, nil
}

// QueryAll 查询所有类型的DNS记录
func (r *LocalResolver) QueryAll(ctx context.Context, domain string) (*QueryResult, error) {
	if domain == "" {
		return nil, fmt.Errorf("domain cannot be empty")
	}

	start := time.Now()
	result := &QueryResult{
		Domain: domain,
	}

	// Query all record types concurrently
	type queryResult struct {
		recordType RecordType
		records    []DNSRecord
		caa        []CAARecord
		err        error
	}

	queries := []RecordType{TypeA, TypeAAAA, TypeCNAME, TypeMX, TypeTXT, TypeNS}
	results := make(chan queryResult, len(queries)+1)

	// Standard DNS queries
	for _, rt := range queries {
		go func(recordType RecordType) {
			records, err := r.Query(ctx, domain, recordType)
			results <- queryResult{recordType: recordType, records: records, err: err}
		}(rt)
	}

	// CAA query
	go func() {
		caa, err := r.QueryCAA(ctx, domain)
		results <- queryResult{recordType: TypeCAA, caa: caa, err: err}
	}()

	// Collect results
	for i := 0; i < len(queries)+1; i++ {
		res := <-results
		if res.err != nil {
			continue // Ignore errors for individual queries
		}

		switch res.recordType {
		case TypeA:
			result.A = res.records
		case TypeAAAA:
			result.AAAA = res.records
		case TypeCNAME:
			result.CNAME = res.records
		case TypeMX:
			result.MX = res.records
		case TypeTXT:
			result.TXT = res.records
		case TypeNS:
			result.NS = res.records
		case TypeCAA:
			result.CAA = res.caa
		}
	}

	result.QueryTime = time.Since(start)
	return result, nil
}

// ResolveCNAMEChain 解析CNAME链
func (r *LocalResolver) ResolveCNAMEChain(ctx context.Context, domain string) ([]string, error) {
	if domain == "" {
		return nil, fmt.Errorf("domain cannot be empty")
	}

	chain := []string{domain}
	seen := make(map[string]bool)
	seen[domain] = true

	currentDomain := domain
	maxDepth := 10 // Prevent infinite loops

	for i := 0; i < maxDepth; i++ {
		cname, err := r.resolver.LookupCNAME(ctx, currentDomain)
		if err != nil {
			// Not a CNAME or error, stop here
			break
		}

		// Remove trailing dot
		cname = strings.TrimSuffix(cname, ".")

		// Check if we've reached the end or a loop
		if cname == currentDomain || seen[cname] {
			break
		}

		chain = append(chain, cname)
		seen[cname] = true
		currentDomain = cname
	}

	return chain, nil
}

// DNSSECState DNSSEC 验证状态
type DNSSECState string

const (
	DNSSECStateSecure        DNSSECState = "secure"        // 已验证（AD=1）
	DNSSECStateInsecure      DNSSECState = "insecure"      // 未启用签名（父区无DS）
	DNSSECStateBogus         DNSSECState = "bogus"         // 签名验证失败
	DNSSECStateIndeterminate DNSSECState = "indeterminate" // 超时或无法完成验证
)

// DNSSECRiskLevel DNSSEC 风险级别
type DNSSECRiskLevel string

const (
	DNSSECRiskNone    DNSSECRiskLevel = "none"    // 无风险（Secure）
	DNSSECRiskWarning DNSSECRiskLevel = "warning" // 警告（Insecure）
	DNSSECRiskDanger  DNSSECRiskLevel = "danger"  // 危险（Bogus/Indeterminate）
)

// DNSSECInfo DNSSEC 检测信息
type DNSSECInfo struct {
	AD        bool            `json:"ad"`         // Authenticated Data 标志
	DO        bool            `json:"do"`         // DNSSEC OK（请求时设置）
	CD        bool            `json:"cd"`         // Checking Disabled（请求时设置）
	TC        bool            `json:"tc"`         // Truncated 标志
	Rcode     string          `json:"rcode"`      // 响应码（NOERROR/SERVFAIL/NXDOMAIN）
	State     DNSSECState     `json:"state"`      // DNSSEC 状态
	RiskLevel DNSSECRiskLevel `json:"risk_level"` // 风险级别
	Message   string          `json:"message,omitempty"` // 状态说明
}
